import ephem
import datetime
from astropy.coordinates import SkyCoord
import astropy.units as u
import numpy as np
from astropy.table import Table
from pandas import DataFrame

class MinorPlanet:
    def __init__(self,timestring,MPCFileName=None,lon=101.181,lat=25.5286,elevation=1950):
        self.MPCName=MPCFileName
        self.timestring=timestring
        self.Catalog=Table()
        self.observatory=ephem.Observer()
        self.observatory.lon=lon
        self.observatory.lat=lat
        self.observatory.date=self.timestring
        self.observatory.elevation=elevation

    def run(self):
        '''
        useage:
        timestring='2019-01-01 02:02:02'
        '''
        infos=[]
        f=open(self.MPCName,'r')
        ephemFormatdata=f.read()
        f.close()
        lef=ephemFormatdata.split('\n')
        #infocat=Table(names=('Name','RA','DEC','MAG'),dtype=('<U100','float32','float32','float32'))
        for i in range(len(lef)-1): #the last lef could be '' which is a zero-length object
            ea=ephem.readdb(lef[i])
            ea.compute(self.observatory)
            infos.append([i,ea.name,ea.mag,np.degrees(ea.a_ra),np.degrees(ea.a_dec)])
            #infocat.add_row([ea.name,np.degrees(ea.a_ra),np.degrees(ea.a_dec),ea.mag])
        infoTab=cc=Table(np.array(infos),names=('ID',"NAME",'MAG','RA','DEC'))
        infoTab['ID']=np.int32(infoTab['ID'])
        infoTab['RA']=np.float32(infoTab['RA'])
        infoTab['DEC']=np.float32(infoTab['DEC'])
        self.Catalog=infoTab
    def region(self,ra,dec,radius): 
        #catalog=self.run(timestring)
        ObjCoor=SkyCoord(ra*u.deg,dec*u.deg)
        MPCoor =SkyCoord(self.Catalog['RA']*u.deg,self.Catalog['DEC']*u.deg)
        pos=np.where(ObjCoor.separation(MPCoor).arcminute<radius)[0]
        return self.Catalog[pos]
    def region2(timestring,ra,dec,radius):
        mpcc=MinorPlanet(timestring)
        mpcc.run()
        ObjCoor=SkyCoord(ra*u.deg,dec*u.deg)
        MPCoor =SkyCoord(mpcc.Catalog['RA']*u.deg,mpcc.Catalog['DEC']*u.deg)
        pos=np.where(ObjCoor.separation(MPCoor).deg<radius)[0]
        return mpcc.Catalog[pos]
    
    def convertFile(inputFileName,outputFileName):
        f=open(inputFileName,'r')
        fdata=f.read()
        f.close()
        ephemFormat=MinorPlanet.convertMpcToXephem(fdata)
        f=open(outputFileName,'w')
        f.write(ephemFormat)
        f.close()
    def convertMpcToXephem(fdata):
        '''
        f=open(mpc_filename,'r')
        fdata=f.read()
        f.close()
        xephemFormat=convertMpcToXephem(fdata)
        '''
        ldata=fdata.split('\n')
        FullText=''
        for i in range(43,len(ldata)):
            lin=ldata[i]
            if len(lin)==0:
                continue
            try:
                res=MinorPlanet.convert_mpcorb_to_xephem(lin).splitlines()[1]
                FullText+=res+'\n'
            except:
                continue
        return FullText
    def _convert_datetime_to_pyephem_date_string(in_datetime):
        return in_datetime.strftime('%Y/%m/%d %H:%M:%S')
    def get_latlon_from_observatory_code(code):
        if type(code) is int:
            code = '%03i' % code
        elif type(code) is str:
            code = code[:3]  # trim any remainder, like @399
        try:
            obscode_filename = _find_cached_file('ObsCodes.html')
            # TODO: add a verbose option to print path to ObsCodes.html
            obscodes = open(obscode_filename, 'r').read().splitlines()
        except:
            raise Error("Problem reading ObsCodes.html file from disk.  \n"
                        "Most likely you need to go download a copy from:  \n"
                        "    http://www.minorplanetcenter.net/iau/lists/ObsCodes.html")
        curobsline = [a for a in obscodes if a.startswith(code)][0]
        output = {'obscode':curobsline[0:3],
                  'longitude':float(curobsline[4:13]),
                  'cos':float(curobsline[13:21]),
                  'sin':float(curobsline[21:30]),
                  'name':curobsline[30:].strip()}
        # From the documentation:
        # "The following list gives the observatory code, longitude (in degrees east of Greenwich) 
        #and the parallax
        # constants (rho cos phi' and rho sin phi', where phi' is the geocentric latitude and rho is 
        #the geocentric
        # distance in earth radii) for each observatory. It is updated nightly."
        output['latitude'] = np.degrees(np.arctan2(output['sin'], output['cos']))
        # Unsure where the following comment came from:
        # geocentric distance in earth radii:
        #   output['sin']/np.sin(np.radians(output['latitude']))
        # NOTE: while ObsCodes.html is clear about being geocentric, it is unclear what pyephem wants, 
        #so blaze ahead
        # TODO: confirm whether pyephem wants geocentric
        return output
    def convert_mpcorb_to_xephem(input):
        """
        convert from, e.g.:
        [MPCORB.DAT](http://www.minorplanetcenter.net/iau/MPCORB/MPCORB.DAT)
        [Format is described in more detail](http://www.minorplanetcenter.org/iau/info/MPOrbitFormat.html)
        Des'n     H     G   Epoch     M        Peri.      Node       Incl.       e            n     \
        a        Reference #Obs #Opp    Arc    rms  Perts   Computer
        # 944 Hidalgo line as of 2013-07-26 is:
        00944   10.77  0.15 K134I 215.40344   56.65077   21.56494   42.54312  0.6617811  0.07172582  \
        5.7370114  0 MPO263352   582  21 1920-2010 0.77 M-v 38h MPCLINUX   0000    (944) Hidalgo   20100222

        to
        # From MPO263352
        944 Hidalgo,e,42.5431,21.5649,56.6508,5.737011,0.0717258,0.66178105,215.4034,04/18.0/2013,\ 
                                                                            2000,H10.77,0.15

        input is a single line of text, output will include a newline character within it (but no newline at end)
        """
        output = '# From ' + input[107:116] + '\n'
        output += input[166:194].strip().replace('(','').replace(')','') + ','
        output += 'e,'
        output += input[59:68].strip() + ','    # i = inclination, degrees
        output += input[48:57].strip() + ','    # O = longitude of ascending node, degrees
        output += input[37:46].strip() + ','    # o = argument of perihelion, degrees
        output += input[92:103].strip() + ','   # a = mean distance (aka semi-major axis), AU
        output += input[80:91].strip() + ','    # n = mean daily motion, degrees per day 
                                                #(computed from a**3/2 if omitted)
        output += input[70:79].strip() + ','    # e = eccentricity, must be < 1
        output += input[26:35].strip() + ','    # M = mean anomaly, i.e., degrees from perihelion
        output += MinorPlanet.unpack_mpc_date(input[20:25].strip()).strftime('%m/%d/%Y') + ','    
                                                # E = epoch date, i.e., time of M
        output += '2000,'   # D = the equinox year, i.e., time of i, O and o  (always J2000.0 in MPCORB.DAT
        output += 'H' + input[8:13].strip() + ','   # First component of magnitude model
        output += input[14:19].strip()   # Second component of magnitude model
        return output
    def unpack_mpc_date(in_packed):
        """
        Convert a MPC packed date format (as described below) to a datetime.date or datetime.datetime object
            http://www.minorplanetcenter.net/iau/info/PackedDates.html

        Copy of the packing definition from the above web page:
            Packed Dates
            Dates of the form YYYYMMDD may be packed into five characters to conserve space.
            The first two digits of the year are packed into a single character in column 1 
            #(I = 18, J = 19, K = 20). Columns 2-3 contain the last two digits of the year. 
            #Column 4 contains the month and column 5 contains the day, coded as detailed below:
               Month     Day      Character         Day      Character
                                 in Col 4 or 5              in Col 4 or 5
               Jan.       1           1             17           H
               Feb.       2           2             18           I
               Mar.       3           3             19           J
               Apr.       4           4             20           K
               May        5           5             21           L
               June       6           6             22           M
               July       7           7             23           N
               Aug.       8           8             24           O
               Sept.      9           9             25           P
               Oct.      10           A             26           Q
               Nov.      11           B             27           R
               Dec.      12           C             28           S
                         13           D             29           T
                         14           E             30           U
                         15           F             31           V
                         16           G

            Examples:

               1996 Jan. 1    = J9611
               1996 Jan. 10   = J961A
               1996 Sept.30   = J969U
               1996 Oct. 1    = J96A1
               2001 Oct. 22   = K01AM
            This system can be extended to dates with non-integral days. The decimal fraction of the \
            #day is simply appended to the five characters defined above.
            Examples:

               1998 Jan. 18.73     = J981I73
               2001 Oct. 22.138303 = K01AM138303
        """
        translate = {}
        for i in range(10):
            translate[str(i)] = i
        for i in range(10,32):
            translate[chr(ord('A') + i - 10)] = i
        if in_packed[0] == 'I':
            year = 1800
        elif in_packed[0] == 'J':
            year = 1900
        elif in_packed[0] == 'K':
            year = 2000
        else:
            raise Error('Unrecognized century code at start of: ' + in_packed)
        year += int(in_packed[1:3])
        month = translate[in_packed[3]]
        day = translate[in_packed[4]]
        if len(in_packed) == 5:
            return datetime.date(year, month, day)
        else:
            decimaldays = float('0.' + in_packed[5:])
            hour = int(decimaldays * 24.)
            minute = int((decimaldays * 24. - hour) * 60.)
            second = int(round(decimaldays * 24. * 60. * 60. - (hour * 3600.) - (minute * 60.)))
            return datetime.datetime(year, month, day, hour, minute, second)
    def pack_mpc_date(in_datetime):
        """
        Convert a datetime.date or datetime.datetime object into the MPC packed date format, as described at:
            http://www.minorplanetcenter.net/iau/info/PackedDates.html

        Copy of the packing definition from the above web page:
            Packed Dates
            Dates of the form YYYYMMDD may be packed into five characters to conserve space.
            The first two digits of the year are packed into a single character in column 1 \
            (I = 18, J = 19, K = 20). Columns 2-3 contain the last two digits of the year. \
            Column 4 contains the month and column 5 contains the day, coded as detailed below:
               Month     Day      Character         Day      Character
                                 in Col 4 or 5              in Col 4 or 5
               Jan.       1           1             17           H
               Feb.       2           2             18           I
               Mar.       3           3             19           J
               Apr.       4           4             20           K
               May        5           5             21           L
               June       6           6             22           M
               July       7           7             23           N
               Aug.       8           8             24           O
               Sept.      9           9             25           P
               Oct.      10           A             26           Q
               Nov.      11           B             27           R
               Dec.      12           C             28           S
                         13           D             29           T
                         14           E             30           U
                         15           F             31           V
                         16           G

            Examples:

               1996 Jan. 1    = J9611
               1996 Jan. 10   = J961A
               1996 Sept.30   = J969U
               1996 Oct. 1    = J96A1
               2001 Oct. 22   = K01AM
            This system can be extended to dates with non-integral days. 
            The decimal fraction of the day is simply appended to the five characters defined above.
            Examples:

               1998 Jan. 18.73     = J981I73
               2001 Oct. 22.138303 = K01AM138303
        """
        if in_datetime.year >= 1800 and in_datetime.year < 1900:
            century = 'I'
        elif in_datetime.year >= 1900 and in_datetime.year < 2000:
            century = 'J'
        elif in_datetime.year >= 2000 and in_datetime.year < 2100:
            century = 'K'
        else:
            raise Error("Year is not within 1800-2099: " + in_datetime.isoformat())
        year = in_datetime.strftime('%y')
        translate = {}
        for i in range(10):
            translate[i] = str(i)
        for i in range(10,32):
            translate[i] = chr(ord('A') + i - 10)
        month = translate[in_datetime.month]
        day = translate[in_datetime.day]
        try:
            decimaldays = ('%7.5f' % ((in_datetime.hour + (in_datetime.minute / 60.) +
                                       (in_datetime.second / 3600.)) / 24.))[2:]
        except:
            decimaldays = ''
        return century + year + month + day + decimaldays
    def minorplanets(in_datetime, observatory_code,
                 max_objects=None,
                 max_magnitude=None, require_magnitude=True,
                 max_zenithdistance_deg=90.0,
                 min_heliocentric_distance_AU=None, max_heliocentric_distance_AU=None,
                 min_topocentric_distance_AU=None, max_topocentric_distance_AU=None):
        """
        in_datetime - datetime.datetime(), e.g. datetime.datetime.utcnow()
        observatory_code - the Code of the observatory in
                           http://www.minorplanetcenter.net/iau/lists/ObsCodes.html
                           can be either string or integer.
        max_objects - default is None, otherwise limits the return to this number
                      of observable objects
        max_magnitude - default is None, otherwise limits return to objects
                        brighter than or equal to this magnitude
                        (as calculated by PyEphem from the MPC data)
                        (TODO: confirm whether this is V-band, R-band,
                        or other...)
        require_magnitude - default is True.  If False and max_magnitude is None,
                            then return all objects, whether PyEphem can calculate
                            a magnitude or not.
        max_zenithdistance_deg - default is 90 degrees (horizon)
        min/max_heliocentric_distance_AU - defaults are None
        min/max_topocentric_distance_AU - defaults are None
        """
        obs_info = MinorPlanet.get_latlon_from_observatory_code(observatory_code)
        obs = ephem.Observer()
        obs.lat = np.radians(obs_info['latitude'])
        obs.lon = np.radians(obs_info['longitude'])
        obs.date = MinorPlanet._convert_datetime_to_pyephem_date_string(in_datetime)
        mpc_filename = MinorPlanet._find_cached_file('MPCORB.DAT')
        if mpc_filename == 'File Not Found':
            raise Error("Problem reading MPCORB.DAT file from disk.  \n"
                        "Most likely you need to go download a copy from:  \n"
                        "    http://www.minorplanetcenter.net/iau/MPCORB/MPCORB.DAT")
        if max_magnitude is not None:
            require_magnitude = True
        matching_objects = []
        with open(mpc_filename) as f:
            in_header = True
            for line in f:
                if in_header is False and len(line) > 1:
                    if (not require_magnitude) or (require_magnitude and (line[8:13] != '     ')):
                        eph = ephem.readdb(convert_mpcorb_to_xephem(line).splitlines()[1])
                        eph.compute(obs)
                        if (max_magnitude is None) or (eph.mag <= max_magnitude):
                            if ((max_zenithdistance_deg is None) or
                                (np.degrees(np.pi/2. - eph.alt) <= max_zenithdistance_deg)):
                                if ((min_heliocentric_distance_AU is None) or
                                    (eph.sun_distance >= min_heliocentric_distance_AU)):
                                    if ((max_heliocentric_distance_AU is None) or
                                        (eph.sun_distance <= max_heliocentric_distance_AU)):
                                        if ((min_topocentric_distance_AU is None) or
                                            (eph.earth_distance >= min_topocentric_distance_AU)):
                                            if ((max_topocentric_distance_AU is None) or
                                                (eph.earth_distance <= max_topocentric_distance_AU)):
                                                matching_objects.append(eph)
                else:
                    if line.startswith('-------------------'):
                        in_header = False
                if max_objects is not None:
                    if len(matching_objects) >= max_objects:
                        break
        name = [a.name for a in matching_objects]
        d = {}
        d['rise_time'] = [a.rise_time.datetime() if a.rise_time is not None else np.nan \
                              for a in matching_objects]
        d['transit_time'] = [a.transit_time.datetime() if a.transit_time is not None else np.nan \
                                 for a in matching_objects]
        d['set_time'] = [a.set_time.datetime() if a.set_time is not None else np.nan for a in matching_objects]
        d['raJ2000_deg'] = [np.degrees(a.a_ra) for a in matching_objects]
        d['decJ2000_deg'] = [np.degrees(a.a_dec) for a in matching_objects]
        d['mag'] = [a.mag for a in matching_objects]
        d['R_AU'] = [a.sun_distance for a in matching_objects]
        d['delta_AU'] = [a.earth_distance for a in matching_objects]
        moon = ephem.Moon()
        moon.compute(obs.date)
        d['O-E-M_deg'] = [np.degrees(ephem.separation(moon, a)) for a in matching_objects]
        output = DataFrame(d, index=name)
        output = output[['rise_time', 'transit_time', 'set_time', 'raJ2000_deg', 'decJ2000_deg',
                         'mag', 'R_AU', 'delta_AU', 'O-E-M_deg']] # re-order columns to something sensible
        return output

if __name__ == "__main__":
    import argparse
    import os
    import sys
    from astropy.io import fits
    from astropy.wcs import WCS
    from astropy.coordinates import SkyCoord
    import astropy.units as u
    import numpy as np
    import pandas as pd
    from loguru import logger

    # ---------------- CLI：保留你的默认值 ----------------
    parser = argparse.ArgumentParser(
        description="在指定 FITS 上用本地 MPC_ephem 数据检索圆锥内小行星，并保存 CSV"
    )
    parser.add_argument(
        "--fits",
        default="/mnt/7b21f1e1-eb25-4cd5-bdb5-06d7d82fa253/Temp/force_photmetry/images/1278305/mb_sc_t1278305-19687_u_20250101132605_194_sciimg.fits",
        help="sciimg FITS 路径",
    )
    parser.add_argument(
        "--mpc-ephem",
        default="./au_data/MPC_ephem.dat",
        help="xephem 目录（每行 ephem.readdb 可解析），若不存在且提供 --mpcorb 将自动生成",
    )
    parser.add_argument(
        "--mpcorb",
        default="./au_data/MPCORB.DAT",
        help="MPCORB.DAT 路径（可选；当 --mpc-ephem 不存在时用于自动转换）",
    )
    parser.add_argument("--time-override", default=None, help="可选：覆盖 FITS 头时间（UTC, 'YYYY-mm-dd HH:MM:SS' 或 ISO）")
    parser.add_argument("--ra-override", type=float, default=None, help="可选：覆盖中心 RA (deg)")
    parser.add_argument("--dec-override", type=float, default=None, help="可选：覆盖中心 Dec (deg)")
    parser.add_argument("--radius", type=float, default=50.0, help="圆锥半径 (arcmin)，默认 50")
    parser.add_argument("--lon", default="100.0313", help="经度（十进制度或 DMS 字符串），默认 100.0313")
    parser.add_argument("--lat", default="26.6974", help="纬度（十进制度或 DMS 字符串），默认 26.6974")
    parser.add_argument("--elev", type=float, default=3200.0, help="海拔（m），默认 3200")
    parser.add_argument("--csv", default="result.csv", help="输出 CSV 路径（默认 result.csv）")
    parser.add_argument("--vmax", type=float, default=None, help="可选：亮度阈值（仅保留 MAG ≤ vmax）")
    args = parser.parse_args()

    # ---------------- 准备 ephem 目录：必要时从 MPCORB 转换 ----------------
    if not os.path.exists(args.mpc_ephem):
        if args.mpcorb and os.path.exists(args.mpcorb):
            logger.info(f"未找到 {args.mpc_ephem}，从 {args.mpcorb} 转换生成 ...")
            with open(args.mpcorb, "r") as f:
                fdata = f.read()
            try:
                xephem_txt = MinorPlanet.convertMpcToXephem(fdata)
                os.makedirs(os.path.dirname(args.mpc_ephem) or ".", exist_ok=True)
                with open(args.mpc_ephem, "w") as f:
                    f.write(xephem_txt)
                logger.info(f"已生成 {args.mpc_ephem}")
            except Exception as e:
                logger.error(f"从 MPCORB 转换失败：{e}")
                sys.exit(1)
        else:
            logger.error(f"找不到 {args.mpc_ephem}，且未提供可用的 --mpcorb 路径。")
            sys.exit(1)

    # ---------------- 从 FITS 取时间与中心坐标（可被 override） ----------------
    try:
        hdr = fits.getheader(args.fits, 0)
    except Exception as e:
        logger.error(f"读取 FITS 失败：{e}")
        sys.exit(1)

    # 时间
    utc = args.time_override or hdr.get("DATE-OBS") or hdr.get("DATE")
    if utc is None:
        logger.error("FITS 头缺少 DATE-OBS/DATE，且未提供 --time-override")
        sys.exit(1)
    utc_for_minorplanet = utc.replace("T", " ")
    logger.info(f"UTC: {utc}")

    # 坐标：优先覆盖；否则 TEL_RA/DEC；否则 CRVAL1/2；最后退回 WCS 图心
    ra0 = args.ra_override if args.ra_override is not None else (hdr.get("TEL_RA") or hdr.get("CRVAL1"))
    dec0 = args.dec_override if args.dec_override is not None else (hdr.get("TEL_DEC") or hdr.get("CRVAL2"))
    if (ra0 is None) or (dec0 is None):
        try:
            w = WCS(hdr)
            nx, ny = hdr.get("NAXIS1"), hdr.get("NAXIS2")
            if (nx is None) or (ny is None) or (not w.is_celestial):
                raise ValueError("无有效 WCS/NAXIS")
            cen = w.pixel_to_world(nx/2, ny/2)
            ra0, dec0 = cen.ra.deg, cen.dec.deg
            logger.warning("未找到 TEL_RA/DEC 与 CRVAL1/2，已使用 WCS 图心作为中心坐标")
        except Exception as e:
            logger.error(f"无法确定中心 RA/Dec：{e}")
            sys.exit(1)
    c_center = SkyCoord(float(ra0)*u.deg, float(dec0)*u.deg, frame="icrs")
    logger.info(f"中心坐标 (deg): RA={c_center.ra.deg:.8f}, Dec={c_center.dec.deg:.8f}")

    # ---------------- 经纬参数：支持十进制度或 DMS ----------------
    def _as_pyephem_coord(x):
        try:
            # 可被 float() 则视为十进制度 → 转弧度
            val = float(x)
            return np.deg2rad(val)
        except Exception:
            # 否则当作 DMS 字符串（如 '+100d01m52.68s'）
            return x

    # ---------------- 调用你的 MinorPlanet ----------------
    try:
        m = MinorPlanet(
            "%s" % utc_for_minorplanet,
            MPCFileName=args.mpc_ephem,
            lon=_as_pyephem_coord(args.lon),
            lat=_as_pyephem_coord(args.lat),
            elevation=args.elev
        )
        m.run()
    except Exception as e:
        logger.error(f"MinorPlanet.run() 失败：{e}")
        sys.exit(1)

    try:
        targetskyarea = m.region(c_center.ra.deg, c_center.dec.deg, args.radius)
    except Exception as e:
        logger.error(f"region() 失败：{e}")
        sys.exit(1)

    logger.info(f"检索到 {len(targetskyarea)} 个目标（半径 {args.radius} arcmin）。")

    # ---------------- 结果整理 + 可选亮度过滤 ----------------
    tab = targetskyarea
    if len(tab) == 0:
        # 仍然写个空 CSV，便于流程可追踪
        pd.DataFrame(columns=["name","number","mag","ra","dec","utc","fits","radius_arcmin"]).to_csv(args.csv, index=False)
        logger.info(f"未检索到目标。已写空文件：{args.csv}")
        sys.exit(0)

    if args.vmax is not None and "MAG" in tab.colnames:
        tab = tab[tab["MAG"] <= float(args.vmax)]
        logger.info(f"亮度过滤（MAG ≤ {args.vmax}）后剩余 {len(tab)} 个目标。")

    rows = []
    for row in tab:  # 列：('ID','NAME','MAG','RA','DEC')
        name_field = str(row["NAME"]).strip()
        # 尝试拆出 number + name
        parts = name_field.split()
        number = None
        name_out = name_field
        if len(parts) >= 2:
            p0 = parts[0].strip("()")
            if p0.isdigit():
                number = p0
                name_out = " ".join(parts[1:])
        elif len(parts) == 1:
            p0 = parts[0].strip("()")
            if p0.isdigit():
                number = p0
        rows.append({
            "name": name_out,
            "number": number,
            "mag": float(row["MAG"]) if row["MAG"] is not None else np.nan,
            "ra": float(row["RA"]),
            "dec": float(row["DEC"]),
            "utc": utc,
            "fits": os.path.basename(args.fits),
            "radius_arcmin": float(args.radius)
        })

    df = pd.DataFrame(rows).sort_values(["mag","ra"], na_position="last")
    # 打印前几行
    logger.info("\n" + df.head(20).to_string(index=False))

    # ---------------- 保存 CSV ----------------
    os.makedirs(os.path.dirname(args.csv) or ".", exist_ok=True)
    df.to_csv(args.csv, index=False)
    logger.info(f"已保存：{args.csv}")

    

        
        